import numpy as np
import rasterio
import os
import scipy.ndimage

def calculate_average_pixel(input_path1, input_path2, output_path):
    # 读取两幅影像
    with rasterio.open(input_path1) as src1, rasterio.open(input_path2) as src2:
        # 获取影像的元数据
        profile = src1.profile
        # 读取影像数据
        data1 = src1.read(1)
        data2 = src2.read(1)

        # 逐像素计算平均值
        average_data = (data1 + data2) / 2.0

    # 将结果保存为新的影像
    with rasterio.open(output_path, 'w', **profile) as dst:
        dst.write(average_data, 1)





def cal_wi2015(input_band2,input_band3, input_band4,input_band5,input_band7,output_path):
    """
    计算wi2015水体指数，输入的tif为Landsat C2 L2 DN值影像，输出水体指数栅格
    （该函数使用过于麻烦，已不建议使用）
    """
        # 读取两幅影像
    with rasterio.open(input_band2) as src2, rasterio.open(input_band3) as src3 \
        ,rasterio.open(input_band4) as src4,rasterio.open(input_band5) as src5 \
        ,rasterio.open(input_band7) as src7:
        # 获取影像的元数据
        profile = src2.profile
        # 读取影像数据，此处读取DN值
        data2 = src2.read(1)
        data3 = src3.read(1)
        data4 = src4.read(1)
        data5 = src5.read(1)
        data7 = src7.read(1)


        #获取每个波段的相对反射率
        # b2=data2*0.00308
        # b3=data3*0.0029
        # b4=data4*0.00345-0.01724
        # b5=data5*0.00235-0.01176
        # b7=data7*0.00364-0.01818

        b2=data2*0.0000275-0.2
        b3=data3*0.0000275-0.2
        b4=data4*0.0000275-0.2
        b5=data5*0.0000275-0.2
        b7=data7*0.0000275-0.2
        

        # 逐像素计算WI2015水体指数
        # wi2015=1.7204+171*data2+3*data3-70*data4-45*data5-71*data7
        wi2015=1.7204+171*b2+3*b3-70*b4-45*b5-71*b7
    # 将结果保存为新的影像
    with rasterio.open(output_path, 'w', **profile) as dst:#使用输入的第一个影像的元数据设置输出影像的元数据，保证它们的空间参考、数据类型等信息一致
        dst.write(wi2015, 1)


def cal_row(input_path,gain,offset,output_path):
    """
    计算输入影像的相对反射率（该函数已弃用）
    """
    with rasterio.open(input_path) as src1:
        # 获取影像的元数据
        profile = src1.profile
        # 读取影像数据
        data1 = src1.read(1)

        # 逐像素计算相对反射率
        output_data=data1*gain+offset

    # 将结果保存为新的影像
    with rasterio.open(output_path, 'w', **profile) as dst:
        dst.write(output_data, 1)





def cal_wi2015_withfolder(input_folder,output_path,thresholded=False):
    """
    输入：input_folder:放置分波段Landsat影像的文件夹，tif文件名必须以Bn（n为波段号）结尾\n
    output_path：tif文件输出路径\n
    thresholded表示是否输出二值化栅格，默认为否
    """
     # 获取文件夹中所有文件
    print("Running...")
    all_files = os.listdir(input_folder)

    # 筛选文件名结尾为B2, B3, B4, B5, B7的文件
    selected_files = [f for f in all_files if f.endswith(('B2.TIF', 'B3.TIF', 'B4.TIF', 'B5.TIF', 'B7.TIF'))]

    if len(selected_files) < 2:
        print("Error: Insufficient input files for calculation.")
        return
    
    input_band2 = os.path.join(input_folder, selected_files[0])
    input_band3 = os.path.join(input_folder, selected_files[1])
    input_band4 = os.path.join(input_folder, selected_files[2])
    input_band5 = os.path.join(input_folder, selected_files[3])
    input_band7 = os.path.join(input_folder, selected_files[4])
    print("Calculating:30%")
    with rasterio.open(input_band2) as src2, rasterio.open(input_band3) as src3 \
        ,rasterio.open(input_band4) as src4,rasterio.open(input_band5) as src5 \
        ,rasterio.open(input_band7) as src7:
        # 获取影像的元数据
        profile = src2.profile
        profile['count']=1
        profile['dtype']='uint16'
        # 读取影像数据，此处读取DN值
        data2 = src2.read(1)
        data3 = src3.read(1)
        data4 = src4.read(1)
        data5 = src5.read(1)
        data7 = src7.read(1)
        print("Calculating:60%")
        b2=data2*0.0000275-0.2
        b3=data3*0.0000275-0.2
        b4=data4*0.0000275-0.2
        b5=data5*0.0000275-0.2
        b7=data7*0.0000275-0.2

        

        # 逐像素计算WI2015水体指数
        # wi2015=1.7204+171*data2+3*data3-70*data4-45*data5-71*data7
        wi2015=1.7204+171*b2+3*b3-70*b4-45*b5-71*b7
        print("Calculating:90%")
    # 将结果保存为新的影像
    if thresholded==False:
        with rasterio.open(output_path, 'w', **profile) as dst:#使用输入的第一个影像的元数据设置输出影像的元数据，保证它们的空间参考、数据类型等信息一致
            dst.write(wi2015, 1)
    else:
        # 阈值分割
        # thresholded_data = np.where(wi2015 > 10000, 0, 1)
        # profile['nodata']=None
        profile['nodata']=None  #重要：否则0值（非水体部分）会被视为NoData
        thresholded_data = np.where(wi2015>0 , 1, 0)
        with rasterio.open(output_path, 'w', **profile) as dst:
            dst.write(thresholded_data, 1)
    print("done")



def cal_wi2015_stackedlayer(input_path,bands,output_path,thresholded=False):
    """
    计算wi2015水体指数，该函数用于读取一个波段合成后的影像，输出水体指数栅格\n
    input_path是输入影像路径，output_path是输出影像路径\n
    bands是一个列表，其中需要包含多波段影像中参与计算WI2015指数的波段序列，
    其顺序要求Landsat的band2, band3, band4, band5, band7\n
    thresholded表示是否输出二值化栅格，默认为否
    """
    print("Running...")
        # 读取影像
    with rasterio.open(input_path) as src:
        # 获取影像的元数据
        profile = src.profile
        profile['count']=1
        # 读取影像数据，此处读取DN值
        data2 = src.read(bands[0])
        data3 = src.read(bands[1])
        data4 = src.read(bands[2])
        data5 = src.read(bands[3])
        data7 = src.read(bands[4])
        print("Calculating:30%")

        b2=data2*0.0000275-0.2
        b3=data3*0.0000275-0.2
        b4=data4*0.0000275-0.2
        b5=data5*0.0000275-0.2
        b7=data7*0.0000275-0.2
        print("Calculating:60%")

        # 逐像素计算WI2015水体指数
        # wi2015=1.7204+171*data2+3*data3-70*data4-45*data5-71*data7
        wi2015=1.7204+171*b2+3*b3-70*b4-45*b5-71*b7
        print("Calculating:90%")
    # 将结果保存为新的影像
    if thresholded==False:
        with rasterio.open(output_path, 'w', **profile) as dst:#使用输入的第一个影像的元数据设置输出影像的元数据，保证它们的空间参考、数据类型等信息一致
            dst.write(wi2015, 1)
    else:
        # 阈值分割
        profile['nodata']=None
        thresholded_data = np.where(wi2015 > 0, 1, 0)
        with rasterio.open(output_path, 'w', **profile) as dst:
            dst.write(thresholded_data, 1)
    print("done")



def threshold_img(input_path,output_path):
    """
    阈值分割
    """
    with rasterio.open(input_path) as src:
        # 获取影像的元数据
        profile = src.profile
        profile['nodata']=None
        profile['count']=1
        data=src.read(1)
        thresholded_data = np.where(data==5, 1, 0)
        with rasterio.open(output_path, 'w', **profile) as dst:
            dst.write(thresholded_data, 1)
    print('done')



def apply_opening(input_path, output_path, iterations=1):
    """
    开运算
    """
    with rasterio.open(input_path, 'r') as src:
        # 读取影像数据
        data = src.read(1)

        # 进行开运算
        opened_data = scipy.ndimage.binary_opening(data, iterations=iterations)

        # 获取元数据
        profile = src.profile

    # 保存结果
    with rasterio.open(output_path, 'w', **profile) as dst:
        dst.write(opened_data, 1)
